package com.lee9213.core.datasource;

import com.lee9213.core.enums.EnumDataSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.*;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 配置主从数据源后，根据选择，返回对应的数据源。多个从库的情况下，会平均的分配从库，用于负载均衡。
 *
 * @author lee9213@163.com
 * @version 1.0
 * @date 2017-11-07 19:54
 */
public class DynamicDataSource extends AbstractRoutingDataSource {
    private final static Logger LOGGER = LoggerFactory.getLogger(DynamicDataSource.class);

    // 主库，只允许有一个
    private DataSource master;
    // 从库，允许有多个
    private List<DataSource> slaves;

    private AtomicLong slaveCount = new AtomicLong();

    private int slaveSize = 0;

    private Map<Object, Object> dataSources = new HashMap<>();

    private static final ThreadLocal<LinkedList<String>> contextHolder = ThreadLocal.withInitial(LinkedList::new);

    /**
     * 初始化数据源
     */
    @Override
    public void afterPropertiesSet() {
        if (master == null) {
            throw new IllegalArgumentException("Property 'master' is required");
        }
        dataSources.put(EnumDataSource.MASTER.getName(), master);
        if (slaves != null && slaves.size() > 0) {
            DataSource slaveDataSource;
            for (int i = 0; i < slaves.size(); i++) {
                slaveDataSource = slaves.get(i);
                dataSources.put(EnumDataSource.SLAVE.getName() + (i + 1), slaveDataSource);
            }
            slaveSize = slaves.size();
        }
        this.setDefaultTargetDataSource(master);
        this.setTargetDataSources(dataSources);
        super.afterPropertiesSet();
    }

    /**
     * 获取数据源
     * 
     * @return
     */
    @Override
    protected Object determineCurrentLookupKey() {
        String dataSource = getDataSource();
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug("current use datasource：{}", dataSource);
        }
        return dataSource;
    }


    /**
     * 获取数据源
     *
     * @return
     */
    private String getDataSource() {
        String dataSourceKey = EnumDataSource.MASTER.getName();
        LinkedList<String> dataSource = contextHolder.get();
        if (dataSource == null) {
            return dataSourceKey;
        }
        String key = dataSource.peekFirst() == null ? EnumDataSource.MASTER.getName() : dataSource.peekFirst();
        if (Objects.equals(EnumDataSource.SLAVE.getName(), key)) {
            if (slaveSize > 1) {
                long c = slaveCount.incrementAndGet();
                c = c % slaveSize;
                dataSourceKey = EnumDataSource.SLAVE.getName() + (c + 1);
            } else {
                dataSourceKey = EnumDataSource.SLAVE.getName() + "1";
            }
        }
        return dataSourceKey;
    }

    /**
     * 设置数据源为Master
     */
    public static void setMasterDataSource() {
        LinkedList<String> dataSource = contextHolder.get();
        dataSource.offerFirst(EnumDataSource.MASTER.getName());
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug("change datasource：{}", dataSource);
        }
    }

    /**
     * 设置数据源为Slave
     */
    public static void setSlaveDataSource() {
        LinkedList<String> dataSource = contextHolder.get();
        dataSource.offerFirst(EnumDataSource.SLAVE.getName());
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug("change datasource：{}", dataSource);
        }
    }

    /**
     * 重置数据源
     */
    public static void reset() {
        LinkedList<String> dataSource = contextHolder.get();
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug("reset datasource {}", dataSource);
        }
        if (dataSource.size() > 0) {
            dataSource.poll();
        }
    }

    /**
     * 清除数据源
     */
    public static void clearDataSource() {
        contextHolder.remove();
    }

    public DataSource getMaster() {
        return master;
    }

    public void setMaster(DataSource master) {
        this.master = master;
    }

    public List<DataSource> getSlaves() {
        return slaves;
    }

    public void setSlaves(List<DataSource> slaves) {
        this.slaves = slaves;
    }
}
